#include <iostream>
#include <opencv2/opencv.hpp>

#include "NvInfer.h"

// 返回结果的结构体
struct Detection {
    cv::Rect box;
    float conf;
    int clas_id;
    std::vector<std::tuple<int, int, float>> keypoints;
};

class YoloDetect {
public:
    // 构造函数初始化模型
    YoloDetect(const std::string &engine_file, const float &conf_threshold, const std::vector<std::string> &class_names,
               const int &point_dim, const std::vector<std::array<int, 2>> &drawLine,
               std::vector<cv::Scalar> &kptColor, std::vector<cv::Scalar> &lineColor,
               const bool &draw_box);

    // 侦测，返回结构体[box,conf,clas_id]
    std::vector<Detection> detect(cv::Mat &draw_img);

    // 析构函数
    ~YoloDetect();

private:
    // 反序列化引擎
    nvinfer1::ICudaEngine *engine = nullptr;
    // 存储中间值的空间
    nvinfer1::IExecutionContext *context = nullptr;

    // 分别创建两个输入和输出的流和内存大小
    void *buffers[2];
    std::vector<int64_t> buffer_size;

    //创建cuda流
    cudaStream_t stream;

    // 缩放图像加黑边
    void ScaleResize(const cv::Mat &image, cv::Mat &resize_image, std::map<std::string, float> &m_resize);

    // BGR --> RGB 标准化 除255
    void ImgNormalize(const cv::Mat &resize_image, cv::Mat &float_image);

    // BGR --> RGB 标准化 减均值除方差
    void ImgNormalize_(const cv::Mat &resize_image, cv::Mat &float_image);

    void ImgHwc2chw(const cv::Mat &float_image, float *blob);

    // 输出处理：置信度阈值筛选，坐标反算，nms
    std::vector<Detection> ScaleNms(const float *output_tensor, const cv::Size &realimg_shape,
                                    const std::map<std::string, float> &m_resize);

    // 画框写标签
    void DrawBox(cv::Mat &draw_img, const std::vector<Detection> &detections);

    // 初始化画框颜色
    void SetColor(const int &num_classes);

    // 置信度阈值
    float confthreshold;
    // 类别容器
    std::vector<std::string> class_names;
    // 每个点的维度
    int point_dim;
    // 关键点连线
    std::vector<std::array<int, 2>> drawLine;
    // 关键点颜色
    std::vector<cv::Scalar> kptColor;
    // 连线颜色
    std::vector<cv::Scalar> lineColor;
    // 是否画框
    bool draw_box;
    // 存储画框颜色
    std::vector<cv::Scalar> clas_colors;

    //获取输入name;
    const char *input_names;
    //获取输出name;
    const char *output_names;

    // 网络需要的输入形状(某个轴可能是动态)，int64_t固定占用8个字节，不受编译器和操作系统影响，为了考虑可移植性
//    std::vector<int64_t> inputtensor_shape;
    // 需要输入图像的 w,h  类型为float
    cv::Size2f input_size2f;
    std::vector<int64_t> output_shape;
    int output_size = 1;  // 输出总形状
};

